#include "qspectrogramdrawing.h"
#include "qaudioreader.h"
#include "qaudiorformat.h"
#include "qfreqalg.h"

#include <QRect>
#include <QPainter>
#include <QImage>

namespace AudioUtils {

	static int getWindowSize(double sampleperpix)
	{
		if (sampleperpix < 256)
		{
			return 256;
		}

		if (sampleperpix < 512)
		{
			return 512;
		}

		if (sampleperpix < 2048)
		{
			return 1024;
		}

		if (sampleperpix < 4096)
		{
			return 512;
		}

		return 256;
	}

	static bool gradient_inited = false;
	static const int gradientSteps = 512;
	static unsigned char gradient_pre[2][2][gradientSteps][3];

	void PreComputeGradient()
	{
		if (!gradient_inited)
		{
			gradient_inited = true;

			for (int selected = 0; selected <= 1; selected++)
			{
				for (int grayscale = 0; grayscale <= 1; grayscale++)
				{
					float r, g, b;

					int i;
					for (i = 0; i < gradientSteps; i++)
					{
						float value = float(i) / gradientSteps;

						if (grayscale)
						{
							r = g = b = 0.84 - 0.84 * value;
						}
						else
						{
							const int gsteps = 4;
							float gradient[gsteps + 1][3] =
							{
								{ float(0), float(0), float(0) },    // lt gray
								{ float(0.1), float(0), float(0.25) },    // lt blue
								{ float(0.5), float(0), float(0.5) },    // violet
								{ float(1.00), float(0.00), float(0.00) },    // red
								{ float(1.00), float(1.00), float(0) }     // white
							};

							int left = int(value * gsteps);
							int right = (left == gsteps ? gsteps : left + 1);

							float rweight = (value * gsteps) - left;
							float lweight = 1.0 - rweight;

							r = (gradient[left][0] * lweight) + (gradient[right][0] * rweight);
							g = (gradient[left][1] * lweight) + (gradient[right][1] * rweight);
							b = (gradient[left][2] * lweight) + (gradient[right][2] * rweight);
						}

						if (selected)
						{
							r *= 0.77f;
							g *= 0.77f;
							b *= 0.885f;
						}

						gradient_pre[selected][grayscale][i][0] = (unsigned char)(255 * r);
						gradient_pre[selected][grayscale][i][1] = (unsigned char)(255 * g);
						gradient_pre[selected][grayscale][i][2] = (unsigned char)(255 * b);
					}
				}
			}
		}
	}

	static void GetColorGradient(
		float value,
		bool selected,
		bool grayscale,
		unsigned char *red,
		unsigned char *green, unsigned char *blue)
	{
		if (!gradient_inited)
			PreComputeGradient();

		if (value >= 0)
		{
			int idx = value * (gradientSteps - 1);
			*red = gradient_pre[selected][grayscale][idx][0];
			*green = gradient_pre[selected][grayscale][idx][1];
			*blue = gradient_pre[selected][grayscale][idx][2];
		}
		else
		{
			*red = 0;
			*green = 0;
			*blue = 0;
		}

	}
	
	QSpectrogramDrawing::QSpectrogramDrawing()
	{
	}

	QSpectrogramDrawing::~QSpectrogramDrawing()
	{
	}

	void QSpectrogramDrawing::draw(
		QPainter& painter, 
		QAudioFormat& format, 
		QAudioReader& reader, 
		const QRect& rect
		)
	{
		quint64 startSample = 0;
		quint64 stopSample = format.samples();
		draw(painter, format, reader, rect, startSample, stopSample);
	}

	void QSpectrogramDrawing::draw(
		QPainter& painter, 
		QAudioFormat& format, 
		QAudioReader& reader, 
		const QRect& rect, 
		quint64 startSample, 
		quint64 stopSample
		)
	{
		if (format.mono())
		{
			draw(painter, format, reader, rect, startSample, stopSample, 0);
		}
		else if (format.stereo())
		{
			QPoint center = rect.center();
			QRect leftRect(rect.topLeft(), QPoint(rect.right(), center.ry()));
			draw(painter, format, reader, leftRect, startSample, stopSample, 0);

			QRect rightRect(QPoint(rect.left(), center.ry()), rect.bottomRight());
			draw(painter, format, reader, rightRect, startSample, stopSample, 1);
		}
	}

	void QSpectrogramDrawing::draw(
		QPainter& painter, 
		QAudioFormat& format, 
		QAudioReader& reader, 
		const QRect& rect, 
		double startSeconds, 
		double stopSeconds
		)
	{
		quint64 startSample = static_cast<quint64>(startSeconds / format.samplerate() + 0.5);
		quint64 stopSample = static_cast<quint64>(stopSeconds / format.samplerate() + 0.5);
		draw(painter, format, reader, rect, startSample, stopSample);
	}

	void QSpectrogramDrawing::draw(
		QPainter& painter, 
		QAudioFormat& format, 
		QAudioReader& reader, 
		const QRect& rect, 
		quint64 startSample, 
		quint64 stopSample, 
		int channel /* 0: left 1: right */
		)
	{
		auto width = rect.width();
		auto height = rect.height();

		if (startSample > stopSample)
		{
			qSwap(startSample, stopSample);
		}

		auto samples = stopSample - startSample;
		auto samplerate = format.samplerate();
		auto sampleperpix = (double)samples / width;

		painter.save();
		painter.fillRect(rect, Qt::black);

		//QPoint origin(rect.left(), rect.center().ry());
		//QTransform transform;
		//transform.translate(rect.left(), rect.bottom());
		//transform.rotate(180, Qt::XAxis);
		//painter.setTransform(transform);

		const int windowSize = getWindowSize(sampleperpix);
		const int half = windowSize / 2;

		float* freq = new float[width * half];
		memset(freq, 0, width * half * sizeof(float));

		quint64* where = new quint64[width + 1];
		for (int x = 0; x < width + 1; x++) 
		{
			// purposely offset the display 1/2 bin to the left (as compared
			// to waveform display to properly center response of the FFT
			where[x] = static_cast<quint64>(x * sampleperpix + 0.5);
		}

		QFreqAlg freqWindow;
		freqWindow.setRate(samplerate);
		freqWindow.setAlg(0);
		freqWindow.setFunc(3);
		freqWindow.setAxis(0);
		freqWindow.setWindowSize(windowSize);

		reader.seek(startSample, SEEK_SET);

		float* buffer = new float[windowSize];
		for (int x = 0; x < width; x++)
		{
			quint64 start = where[x];
			quint64 len = windowSize;

			if (start <= 0 || start >= samples)
			{
				for (quint64 i = 0; i < (quint64)half; i++)
				{
					freq[half * x + i] = 0;
				}
			}
			else
			{
				float* adj = buffer;
				start -= windowSize >> 1;

				if (start < 0) 
				{
					for (quint64 i = start; i < 0; i++)
					{
						*adj++ = 0;
					}
						
					len += start;
					start = 0;
				}

				if (start + len > samples)
				{
					int newlen = samples - start;
					for (quint64 i = newlen; i < len; i++)
					{
						adj[i] = 0;
					}

					len = newlen;
				}

				if (len > 0)
				{
					reader.seek(start, SEEK_SET);
					auto count = reader.readTrack(adj, len, channel);
					for (quint64 i = count; i < len; i++)
					{
						adj[i] = 0;
					}
				}

				freqWindow.setData(buffer, windowSize);
				freqWindow.recalc();
				freqWindow.getFreqDb(&freq[x * half], half);
				// ComputeSpectrumUsingRealFFTf(buffer, hFFT, mWindow, mWindowSize, &mSpecCache->freq[half * x]);
			}
		}
		delete[] where;
		delete[] buffer;

		int ifreq = lrint(samplerate / 2);
		int maxFreq = ifreq;
		int minFreq = 0;

		int minSamples = int((double)minFreq * (double)windowSize / samplerate + 0.5);   // units are fft bins
		int maxSamples = int((double)maxFreq * (double)windowSize / samplerate + 0.5);
		int temp = sizeof(int);
		float binPerPx = float(maxSamples - minSamples) / float(height);
		int range = 80;
		int gain = 20;

		int x = 0;
		quint64 w1 = static_cast<quint64>(x * samplerate + 0.5);

		const float f = samplerate / 2.0f / half;
		const float	lmin = logf(float(minFreq));
		const float	lmax = logf(float(maxFreq));
		const float	scale = lmax - lmin;

		QImage image(width, height, QImage::Format_RGB32);
		for (int x = 0; x < width; x++)
		{
			quint64 w0 = w1;
			w1 = static_cast<quint64>((x + 1) * samplerate + 0.5);
			for (int y = 0; y < height; y++)
			{
				bool selflag = false;
				unsigned char rv, gv, bv;
				float value;

				float bin0 = float(y) * binPerPx + minSamples;
				float bin1 = float(y + 1) * binPerPx + minSamples;

				if (int(bin1) == int(bin0))
				{
					value = freq[half * x + int(bin0)];
				}
				else
				{
					float binwidth = bin1 - bin0;
					value = freq[half * x + int(bin0)] * (1.f - bin0 + (int)bin0);

					bin0 = 1 + int(bin0);
					while (bin0 < int(bin1))
					{
						value += freq[half * x + int(bin0)];
						bin0 += 1.0;
					}

					// Do not reference past end of freq array.
					if (int(bin1) >= half)
					{
						bin1 -= 1.0;
					}

					value += freq[half * x + int(bin1)] * (bin1 - int(bin1));
					value /= binwidth;
				}

				// Last step converts dB to a 0.0-1.0 range
				value = (value + range + gain) / (double)range;
				if (value > 1.0)
				{
					value = float(1.0);
				}

				if (value < 0.0)
				{
					value = float(0.0);
				}

				GetColorGradient(value, selflag, false, &rv, &gv, &bv);
				image.setPixelColor(x, height - y - 1, qRgb(rv, gv, bv));
			}
		}

		delete[] freq;
		painter.drawImage(rect.left(), rect.top(), image);
		painter.restore();
	}

}